// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <boost/random/uniform_int_distribution.hpp>
#include <string>

#include <pollux/plan/input_generator.h>

using namespace kumo::pollux::plan;

namespace kumo::pollux::exec::test {
    class ApproxPercentileInputGenerator : public InputGenerator {
    public:
        std::vector<VectorPtr> generate(
            const std::vector<TypePtr> &types,
            VectorFuzzer &fuzzer,
            FuzzerGenerator &rng,
            memory::MemoryPool *pool) override {
            // The arguments are: x, [w], percentile(s), [accuracy].
            //
            // First argument is always 'x'. If second argument's type is BIGINT, then
            // it is 'w'. Otherwise, it is percentile(x).

            const auto size = fuzzer.getOptions().vectorSize;

            std::vector<VectorPtr> inputs;
            inputs.reserve(types.size());
            inputs.push_back(fuzzer.fuzz(types[0]));

            if (types[1]->is_bigint()) {
                pollux::VectorMaker vectorMaker{pool};
                auto weight = vectorMaker.flat_vector<int64_t>(size, [&](auto /*row*/) {
                    return boost::random::uniform_int_distribution<int64_t>(1, 1'000)(rng);
                });

                inputs.push_back(weight);
            }

            const int percentileTypeIndex = types[1]->is_bigint() ? 2 : 1;
            const TypePtr &percentileType = types[percentileTypeIndex];
            if (percentileType->is_double()) {
                if (!percentile_.has_value()) {
                    percentile_ = pickPercentile(fuzzer, rng);
                }

                inputs.push_back(BaseVector::create_constant(
                    DOUBLE(), percentile_.value(), size, pool));
            } else {
                POLLUX_CHECK(percentileType->is_array());
                POLLUX_CHECK(percentileType->childAt(0)->is_double());

                if (percentiles_.empty()) {
                    percentiles_.push_back(pickPercentile(fuzzer, rng));
                    percentiles_.push_back(pickPercentile(fuzzer, rng));
                    percentiles_.push_back(pickPercentile(fuzzer, rng));
                }

                auto array_vector =
                        BaseVector::create<ArrayVector>(ARRAY(DOUBLE()), 1, pool);
                auto elementsVector = array_vector->elements()->as_flat_vector<double>();
                elementsVector->resize(percentiles_.size());
                for (auto i = 0; i < percentiles_.size(); ++i) {
                    elementsVector->set(i, percentiles_[i]);
                }
                array_vector->setOffsetAndSize(0, 0, percentiles_.size());

                inputs.push_back(BaseVector::wrap_in_constant(size, 0, array_vector));
            }

            if (types.size() > percentileTypeIndex + 1) {
                // Last argument is 'accuracy'.
                POLLUX_CHECK(types.back()->is_double());
                if (!accuracy_.has_value()) {
                    accuracy_ = boost::random::uniform_01<double>()(rng);
                }

                inputs.push_back(
                    BaseVector::create_constant(DOUBLE(), accuracy_.value(), size, pool));
            }

            return inputs;
        }

        void reset() override {
            percentile_.reset();
            percentiles_.clear();
            accuracy_.reset();
        }

    private:
        double pickPercentile(VectorFuzzer &fuzzer, FuzzerGenerator &rng) {
            // 10% of the times generate random value in [0, 1] range.
            // 90% of the times use one of the common values.
            if (fuzzer.coinToss(0.1)) {
                return boost::random::uniform_01<double>()(rng);
            }

            static const std::vector<double> kPercentiles = {
                0.1, 0.25, 0.5, 0.75, 0.90, 0.95, 0.99, 0.999, 0.9999
            };

            const auto index =
                    boost::random::uniform_int_distribution<uint32_t>()(rng) %
                    kPercentiles.size();

            return kPercentiles[index];
        }

        std::optional<double> percentile_;
        std::vector<double> percentiles_;
        std::optional<double> accuracy_;
    };
} // namespace kumo::pollux::exec::test
